package org.apache.spark.ml.feature

import org.apache.hadoop.fs.Path
import org.apache.spark.SparkException
import org.apache.spark.ml.attribute.NominalAttribute
import org.apache.spark.ml.feature.WordIndexerModel.WordIndexerModelWriter
import org.apache.spark.ml.param.ParamMap
import org.apache.spark.ml.util._
import org.apache.spark.ml.{Estimator, Model}
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types._
import org.apache.spark.sql.{DataFrame, Dataset}
import org.apache.spark.util.collection.OpenHashMap

/**
  * @Author songhaicheng
  * @Date 2019/2/19 14:07
  * @Description
  * @Reviewer
  */
private[feature] trait WordIndexerBase extends StringIndexerBase {

  override protected def validateAndTransformSchema(schema: StructType): StructType = {
    val inputColName = $(inputCol)
    val inputDataType = schema(inputColName).dataType
    require(inputDataType == ArrayType(StringType),
      s"The input column $inputColName must be array type of string, " +
        s"but got $inputDataType.")
    val inputFields = schema.fields
    val outputColName = $(outputCol)
    require(inputFields.forall(_.name != outputColName),
      s"Output column $outputColName already exists.")
    val attr = NominalAttribute.defaultAttr.withName($(outputCol))
    val outputFields = inputFields :+ attr.toStructField()
    StructType(outputFields)
  }

}

class WordIndexer (
                    override val uid: String) extends Estimator[WordIndexerModel]
  with WordIndexerBase with DefaultParamsWritable {

  def this() = this(Identifiable.randomUID("wordIdx"))

  def setHandleInvalid(value: String): this.type = set(handleInvalid, value)

  def setInputCol(value: String): this.type = set(inputCol, value)

  def setOutputCol(value: String): this.type = set(outputCol, value)

  override def fit(dataset: Dataset[_]): WordIndexerModel = {
    transformSchema(dataset.schema, logging = true)
    val counts = dataset.na.drop(Array($(inputCol))).select(col($(inputCol)))
      .rdd
      .flatMap(_.getAs[Seq[String]](0))
      .countByValue()
    val labels = counts.toSeq.sortBy(-_._2).map(_._1).toArray
    copyValues(new WordIndexerModel(uid, labels).setParent(this))
  }

  override def copy(extra: ParamMap): WordIndexer = defaultCopy(extra)

  override def transformSchema(schema: StructType): StructType = validateAndTransformSchema(schema)
}


object WordIndexer extends DefaultParamsReadable[WordIndexer] {
  private[feature] val SKIP_INVALID: String = "skip"
  private[feature] val ERROR_INVALID: String = "error"
  private[feature] val KEEP_INVALID: String = "keep"
  private[feature] val supportedHandleInvalids: Array[String] =
    Array(SKIP_INVALID, ERROR_INVALID, KEEP_INVALID)

  override def load(path: String): WordIndexer = super.load(path)
}


class WordIndexerModel (
                         override val uid: String,
                         val labels: Array[String])
  extends Model[WordIndexerModel] with WordIndexerBase with MLWritable {

  def this(labels: Array[String]) = this(Identifiable.randomUID("wordIdx"), labels)

  private val labelToIndex: OpenHashMap[String, Double] = {
    val n = labels.length
    val map = new OpenHashMap[String, Double](n)
    var i = 0
    while (i < n) {
      map.update(labels(i), i)
      i += 1
    }
    map
  }

  def setHandleInvalid(value: String): this.type = set(handleInvalid, value)

  def setInputCol(value: String): this.type = set(inputCol, value)

  def setOutputCol(value: String): this.type = set(outputCol, value)

  override def copy(extra: ParamMap): WordIndexerModel = {
    val copied = new WordIndexerModel(uid, labels)
    copyValues(copied, extra).setParent(parent)
  }

  override def write: WordIndexerModelWriter = new WordIndexerModelWriter(this)

  override def transform(dataset: Dataset[_]): DataFrame = {
    if (!dataset.schema.fieldNames.contains($(inputCol))) {
      logInfo(s"Input column ${$(inputCol)} does not exist during transformation. " +
        "Skip WordIndexerModel.")
      return dataset.toDF
    }
    transformSchema(dataset.schema, logging = true)

    val filteredLabels = getHandleInvalid match {
      case WordIndexer.KEEP_INVALID => labels :+ "__unknown"
      case _ => labels
    }

    val metadata = NominalAttribute.defaultAttr
      .withName($(outputCol)).withValues(filteredLabels).toMetadata()

    // If we are skipping invalid records, filter them out.
    val (filteredDataset, keepInvalid) = getHandleInvalid match {
      case WordIndexer.SKIP_INVALID =>
        val filterer = udf { words: Seq[String] =>
          words.filter(labelToIndex.contains)
        }
        (dataset.na.drop(Array($(inputCol))).withColumn($(inputCol), filterer(dataset($(inputCol)))), false)
      case _ => (dataset, getHandleInvalid == WordIndexer.KEEP_INVALID)
    }

    val indexer = udf { words: Seq[String] =>
      if (words == null) {
        if (keepInvalid) {
          Seq(labels.length.toDouble)
        } else {
          throw new SparkException("WordIndexer encountered NULL value. To handle or skip " +
            "NULLS, try setting WordIndexer.handleInvalid.")
        }
      } else {
        val vec = new Array[Double](words.length)
        for (i <- words.indices) {
          if (labelToIndex.contains(words(i))) {
            vec.update(i, labelToIndex(words(i)))
          } else if (keepInvalid) {
            vec.update(i, labels.length.toDouble)
          } else {
            throw new SparkException(s"Unseen word: ${words(i)}.  To handle unseen words, " +
              s"set Param handleInvalid to ${WordIndexer.KEEP_INVALID}.")
          }
        }
        vec.toSeq
      }
    }

    filteredDataset.select(col("*"), indexer(dataset($(inputCol))).as($(outputCol), metadata))
  }

  override def transformSchema(schema: StructType): StructType = {
    if (schema.fieldNames.contains($(inputCol))) {
      validateAndTransformSchema(schema)
    } else {
      schema
    }
  }
}


object WordIndexerModel extends MLReadable[WordIndexerModel] {

  private[WordIndexerModel]
  class WordIndexerModelWriter(instance: WordIndexerModel) extends MLWriter {

    private case class Data(labels: Array[String])

    override protected def saveImpl(path: String): Unit = {
      DefaultParamsWriter.saveMetadata(instance, path, sc)
      val data = Data(instance.labels)
      val dataPath = new Path(path, "data").toString
      sparkSession.createDataFrame(Seq(data)).repartition(1).write.parquet(dataPath)
    }
  }

  private class WordIndexerModelReader extends MLReader[WordIndexerModel] {

    private val className = classOf[WordIndexerModel].getName

    override def load(path: String): WordIndexerModel = {
      val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
      val dataPath = new Path(path, "data").toString
      val data = sparkSession.read.parquet(dataPath)
        .select("labels")
        .head()
      val labels = data.getAs[Seq[String]](0).toArray
      val model = new WordIndexerModel(metadata.uid, labels)
      DefaultParamsReader.getAndSetParams(model, metadata)
      model
    }
  }

  override def read: MLReader[WordIndexerModel] = new WordIndexerModelReader

  override def load(path: String): WordIndexerModel = super.load(path)
}
